package cn.jzk.aitest.Controller;

import cn.jzk.aitest.DTO.MessageDTO;
import cn.jzk.aitest.Service.ChatHistoryService;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.function.Function;

/**
 * 会话管理
 * @author jiangzk
 * 2025/4/21  上午11:39
 */
@RestController
@RequestMapping("/ai")
public class ChatController {

    @Resource
    private ChatClient chatClient;

    @Resource
    private ChatHistoryService chatHistoryService;

    @Resource
    private ChatMemory chatMemory;

    public static final String CHAT_MEMORY_CONVERSATION_ID_KEY = AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;

    public static final String CHAT_ID_KEY = "chat_id_key";


    /**
     * 模型对话接口
     */
    // 由于以流式方式调用模型，展示的数据会出现乱码，需要进行规定编码格式
    @GetMapping(value = "/chat", produces = "text/html; charset=utf-8")
    public Flux<String> chat(String msg, String chatId) {

        // 1.添加会话信息
        chatHistoryService.save(CHAT_ID_KEY, chatId);

        // 2.调用模型
        return chatClient.prompt()
                .user(msg)
                // 2.1 配置会话记忆功能
                .advisors(advisorSpec -> advisorSpec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                .stream()
                .content();
    }


    /**
     * 查询会话记录详情（这里的 type 可根据需求进行更改，chatId 建议保留以作为会话的唯一ID）
     */
    @GetMapping("/{type}/{chatId}")
    public List<MessageDTO> getChatHistory(@PathVariable("type") String type, @PathVariable("chatId") String chatId) {
        // 这里的展示记录条数最大值规定，直接使用的 int 型最大值，具体自己根据需求规定
        List<Message> messages = chatMemory.get(chatId, Integer.MAX_VALUE);
        if (messages == null) {
            return List.of();
        }
        // 进行封装返回类型
        return messages.stream().map(new Function<Message, MessageDTO>() {
            @Override
            public MessageDTO apply(Message message) {
                return new MessageDTO(message);
            }
        }).toList();
    }
}
